import streamlit as st
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import r2_score
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import pandas as pd

# 设置页面配置
st.set_page_config(
    page_title="视频收入预测系统",
    page_icon="🎬",
    layout="wide"
)

# 设置随机种子以确保可重复性
torch.manual_seed(42)
np.random.seed(42)


# 1. 生成随机数据
def generate_data(num_samples=20):
    plays = np.random.randint(90, 101, size=num_samples)  # 播放量 90-100
    likes = []
    incomes = []

    for play in plays:
        # 点赞数为播放量的1/5，加上±5%的误差
        base_likes = play / 5
        error = np.random.uniform(-0.05, 0.05)  # ±5%误差
        like = base_likes * (1 + error)
        likes.append(like)

        # 收入公式
        income = 0.1 * play + 0.3 * like
        incomes.append(income)

    return np.array(plays), np.array(likes), np.array(incomes)


# 3. 创建神经网络模型
class VideoIncomePredictor(nn.Module):
    def __init__(self, hidden_layers=[64, 32, 16]):
        super(VideoIncomePredictor, self).__init__()
        layers = []
        input_size = 2

        for hidden_size in hidden_layers:
            layers.append(nn.Linear(input_size, hidden_size))
            layers.append(nn.ReLU())
            input_size = hidden_size

        layers.append(nn.Linear(input_size, 1))

        self.network = nn.Sequential(*layers)

    def forward(self, x):
        return self.network(x)


# 初始化session state
if 'trained' not in st.session_state:
    st.session_state.trained = False
if 'model' not in st.session_state:
    st.session_state.model = None
if 'data' not in st.session_state:
    st.session_state.data = None
if 'scalers' not in st.session_state:
    st.session_state.scalers = None

# 页面标题
st.title("🎬 视频收入预测系统")
st.markdown("通过播放量和点赞数预测视频收入")

# 侧边栏配置
with st.sidebar:
    st.header("⚙️ 配置参数")

    num_samples = st.slider("样本数量", 10, 100, 20)
    epochs = st.slider("训练轮次", 100, 5000, 1000)
    learning_rate = st.slider("学习率", 0.001, 0.1, 0.01, 0.001)

    if st.button("🔄 生成并训练模型", type="primary"):
        with st.spinner("正在生成数据并训练模型..."):
            # 生成数据
            plays, likes, incomes = generate_data(num_samples)
            st.session_state.data = (plays, likes, incomes)

            # 数据预处理
            X = np.column_stack((plays, likes))
            y = incomes.reshape(-1, 1)

            # 特征归一化
            scaler_X = StandardScaler()
            X_scaled = scaler_X.fit_transform(X)

            # 目标值归一化
            scaler_y = StandardScaler()
            y_scaled = scaler_y.fit_transform(y)

            # 转换为PyTorch张量
            X_tensor = torch.FloatTensor(X_scaled)
            y_tensor = torch.FloatTensor(y_scaled)

            # 创建和训练模型
            model = VideoIncomePredictor()
            criterion = nn.MSELoss()
            optimizer = optim.Adam(model.parameters(), lr=learning_rate)

            losses = []
            model.train()

            progress_bar = st.progress(0)
            status_text = st.empty()

            for epoch in range(epochs):
                # 前向传播
                predictions = model(X_tensor)
                loss = criterion(predictions, y_tensor)

                # 反向传播
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                losses.append(loss.item())

                # 更新进度
                if (epoch + 1) % 100 == 0:
                    progress = (epoch + 1) / epochs
                    progress_bar.progress(progress)
                    status_text.text(f"训练中... {epoch + 1}/{epochs} 轮次")

            # 保存模型和数据
            st.session_state.model = model
            st.session_state.scalers = (scaler_X, scaler_y)
            st.session_state.losses = losses
            st.session_state.trained = True

            # 评估模型
            model.eval()
            with torch.no_grad():
                test_predictions = model(X_tensor)
                test_predictions_original = scaler_y.inverse_transform(test_predictions.numpy())
                r2 = r2_score(incomes, test_predictions_original.flatten())
                st.session_state.r2 = r2

            st.success("模型训练完成！")

# 主内容区
if st.session_state.trained:
    plays, likes, incomes = st.session_state.data
    model = st.session_state.model
    scaler_X, scaler_y = st.session_state.scalers
    losses = st.session_state.losses

    # 创建标签页
    tab1, tab2, tab3, tab4 = st.tabs(["📊 数据概览", "📈 训练结果", "🔮 收入预测", "📋 模型信息"])

    with tab1:
        st.header("数据集概览")

        # 创建数据表格
        df = pd.DataFrame({
            '播放量': plays,
            '点赞数': likes,
            '实际收入': incomes
        })
        st.dataframe(df, use_container_width=True)

        # 数据统计
        col1, col2, col3 = st.columns(3)
        with col1:
            st.metric("平均播放量", f"{np.mean(plays):.1f}")
        with col2:
            st.metric("平均点赞数", f"{np.mean(likes):.2f}")
        with col3:
            st.metric("平均收入", f"{np.mean(incomes):.2f}")

        # 散点图
        fig = go.Figure()
        fig.add_trace(go.Scatter3d(
            x=plays, y=likes, z=incomes,
            mode='markers',
            marker=dict(size=8, color=incomes, colorscale='Viridis', opacity=0.8),
            name='数据点'
        ))

        fig.update_layout(
            title='视频数据分布 (播放量 vs 点赞数 vs 收入)',
            scene=dict(
                xaxis_title='播放量',
                yaxis_title='点赞数',
                zaxis_title='收入'
            ),
            height=600
        )
        st.plotly_chart(fig, use_container_width=True)

    with tab2:
        st.header("训练结果分析")

        col1, col2 = st.columns(2)
        with col1:
            # 损失曲线
            fig_loss = go.Figure()
            fig_loss.add_trace(go.Scatter(
                x=list(range(len(losses))),
                y=losses,
                mode='lines',
                name='训练损失'
            ))
            fig_loss.update_layout(
                title='训练损失曲线',
                xaxis_title='Epoch',
                yaxis_title='MSE Loss',
                yaxis_type='log',
                height=400
            )
            st.plotly_chart(fig_loss, use_container_width=True)

        with col2:
            # 模型性能指标
            st.metric("R²分数", f"{st.session_state.r2:.4f}")
            st.metric("最终损失", f"{losses[-1]:.6f}")

            # 真实vs预测散点图
            model.eval()
            with torch.no_grad():
                X_scaled = scaler_X.transform(np.column_stack((plays, likes)))
                predictions_scaled = model(torch.FloatTensor(X_scaled)).numpy()
                predictions = scaler_y.inverse_transform(predictions_scaled)

            fig_scatter = go.Figure()
            fig_scatter.add_trace(go.Scatter(
                x=incomes,
                y=predictions.flatten(),
                mode='markers',
                name='预测值'
            ))
            fig_scatter.add_trace(go.Scatter(
                x=[min(incomes), max(incomes)],
                y=[min(incomes), max(incomes)],
                mode='lines',
                name='理想线',
                line=dict(dash='dash')
            ))
            fig_scatter.update_layout(
                title='真实值 vs 预测值',
                xaxis_title='真实收入',
                yaxis_title='预测收入',
                height=400
            )
            st.plotly_chart(fig_scatter, use_container_width=True)

    with tab3:
        st.header("收入预测")

        col1, col2 = st.columns(2)

        with col1:
            st.subheader("输入特征")
            play_count = st.slider("播放量", 90, 100, 95)
            like_count = st.slider("点赞数", 15.0, 25.0, 19.0, 0.1)

            if st.button("预测收入", type="primary"):
                # 准备输入数据
                input_data = np.array([[play_count, like_count]])
                input_scaled = scaler_X.transform(input_data)

                # 预测
                model.eval()
                with torch.no_grad():
                    prediction_scaled = model(torch.FloatTensor(input_scaled)).numpy()
                    prediction = scaler_y.inverse_transform(prediction_scaled)[0][0]

                # 理论收入
                theoretical_income = 0.1 * play_count + 0.3 * like_count

                st.session_state.prediction = prediction
                st.session_state.theoretical = theoretical_income

        with col2:
            if 'prediction' in st.session_state:
                st.subheader("预测结果")

                col21, col22 = st.columns(2)
                with col21:
                    st.metric("预测收入", f"{st.session_state.prediction:.2f}")
                with col22:
                    st.metric("理论收入", f"{st.session_state.theoretical:.2f}")

                # 误差分析
                error = abs(st.session_state.prediction - st.session_state.theoretical)
                error_percent = (error / st.session_state.theoretical) * 100

                st.progress(min(int(error_percent), 100))
                st.write(f"误差: {error:.2f} ({error_percent:.1f}%)")

                # 特征重要性
                st.subheader("特征贡献度")
                play_contribution = 0.1 * play_count
                like_contribution = 0.3 * like_count
                total = play_contribution + like_contribution

                fig_pie = go.Figure(data=[go.Pie(
                    labels=['播放量贡献', '点赞数贡献'],
                    values=[play_contribution, like_contribution],
                    hole=0.4
                )])
                fig_pie.update_layout(height=300)
                st.plotly_chart(fig_pie, use_container_width=True)

    with tab4:
        st.header("模型架构信息")

        # 显示模型结构
        st.code(str(model), language='python')

        # 模型参数统计
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

        col41, col42 = st.columns(2)
        with col41:
            st.metric("总参数数量", f"{total_params:,}")
        with col42:
            st.metric("可训练参数", f"{trainable_params:,}")

        # 特征重要性
        st.subheader("特征标准化参数")
        st.write(f"播放量: 均值={scaler_X.mean_[0]:.2f}, 标准差={scaler_X.scale_[0]:.2f}")
        st.write(f"点赞数: 均值={scaler_X.mean_[1]:.2f}, 标准差={scaler_X.scale_[1]:.2f}")

else:
    st.info("👈 请在侧边栏配置参数并训练模型")

    # 显示示例数据
    st.header("示例数据格式")
    example_plays, example_likes, example_incomes = generate_data(5)

    example_df = pd.DataFrame({
        '播放量': example_plays,
        '点赞数': example_likes,
        '收入': example_incomes
    })
    st.dataframe(example_df)

    st.markdown("""
    ### 收入计算公式：
    `收入 = 0.1 × 播放量 + 0.3 × 点赞数`

    ### 数据生成规则：
    - 播放量：90-100之间的随机整数
    - 点赞数：播放量的1/5 ±5%误差
    """)

# 添加页脚
st.markdown("---")
st.caption("🎬 视频收入预测系统 - 基于PyTorch和Streamlit")